package middleware

import (
	"fiberx/internal/global/response"
	"sync"

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

// userLimiters 用于存储每个用户的速率限制器
var userLimiters = make(map[string]*rate.Limiter)
var mu sync.Mutex

// getUserLimiter 获取或创建用户的速率限制器
func getUserLimiter(userID string) *rate.Limiter {
	mu.Lock()
	defer mu.Unlock()

	limiter, exists := userLimiters[userID]
	if !exists {
		// limiter = rate.NewLimiter(rate.Limit(1), 5)
		// 1 表示每秒生成 1 个令牌，5 表示令牌桶的容量为 5。
		// 假设系统处于空闲状态，令牌桶已经被填满，此时有 5 个请求同时到来。由于令牌桶中有 5 个令牌，这 5 个请求都可以立即被处理，因为每个请求只需要消耗一个令牌。处理完这 5 个请求后，令牌桶为空，后续的请求就需要等待新的令牌生成。

		// 每个用户每秒允许 5 次请求
		limiter = rate.NewLimiter(rate.Limit(5), 10)
		userLimiters[userID] = limiter
	}
	return limiter
}

// RateLimitByUser 按用户限流的中间件
func RateLimitByUser() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 这里以 IP 地址作为用户标识，实际应用中可根据需求修改
		var clientID string
		if client_id, ok := c.Get("clientid"); ok {
			clientID = client_id.(string)
		}
		limiter := getUserLimiter(clientID)

		// 尝试从令牌桶中获取一个令牌
		if !limiter.Allow() {
			c.Set("code", response.ERROR_REQUEST_TOO_FAST)
			c.Abort()
			return
		}
		// 若有足够的令牌，继续处理请求
		c.Next()
	}
}
